/*
MIT License

Copyright (c) 2007 Ralf Kornmann

All right reserved.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 
*/

#pragma once

class CFunctionHook;
class CPlugin;

class CFunction
{
	public:
		CFunction (LPCSTR strProc);

		const std::string& GetName () const;

		void AddHook (CFunctionHook* pFunction);
		void RemoveHook (CFunctionHook* pFunction);

		void UpdateState (HMODULE hModule);

	protected:
		std::string m_strName;

		FARPROC m_Proc;

		typedef std::vector<CFunctionHook*> TFunctionHookVector;
		TFunctionHookVector m_Functions;
};

class CModule
{
	public:
		CModule (LPCWSTR strModule);
		~CModule ();

		const std::wstring& GetName () const;

		CFunction* GetFunction (LPCSTR strProc);

		void UpdateState ();

	protected:
		std::wstring m_strName;
		HMODULE m_hModule;

		typedef std::vector<CFunction*> TFunctionVector;
		TFunctionVector m_Functions;
};

class CHost : public IUAHHost
{
	public:
		CHost ();
		virtual ~CHost ();

        virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppvObject);
        virtual ULONG STDMETHODCALLTYPE AddRef( void);
        virtual ULONG STDMETHODCALLTYPE Release( void);

		virtual HRESULT STDMETHODCALLTYPE CreateFunctionHook(LPCWSTR lpModuleName, LPCSTR lpProcName, FARPROC newFunction, DWORD Flags, IUAHHockedFunction** ppFunction);
		virtual HRESULT STDMETHODCALLTYPE LoadPlugin (LPCWSTR lpPlugin, LPCWSTR lpConfig);
		virtual HRESULT STDMETHODCALLTYPE LoadPlugin (LPCWSTR lpPlugin, XML::IXMLDOMElement* configElement);
		virtual HRESULT STDMETHODCALLTYPE LoadPlugins (XML::IXMLDOMElement* configElement);
		virtual HRESULT STDMETHODCALLTYPE LoadPlugins (LPCWSTR lpConfig);
		virtual HRESULT STDMETHODCALLTYPE LoadPluginsFromFile (LPCWSTR lpFile);

		BOOL GetThreadContext (HANDLE hThread, LPCONTEXT lpContext);
		BOOL SetThreadContext (HANDLE hThread, const CONTEXT* lpContext);
		BOOL VirtualProtect (LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect);
		BOOL WriteProcessMemory (HANDLE hProcess, LPVOID lpBaseAddress, LPCVOID lpBuffer, SIZE_T nSize, SIZE_T* lpNumberOfBytesWritten);

		FARPROC GetPatchFunctionBlock ();
		void RemoveMonitor (CFunctionHook* pFunction);

		void UpdateState ();

	protected:
		LONG m_RefCount;

		typedef BOOL (WINAPI *GetThreadContextFunc) (HANDLE hThread, LPCONTEXT lpContext);
		typedef BOOL (WINAPI *SetThreadContextFunc) (HANDLE hThread, const CONTEXT* lpContext);
		typedef BOOL (WINAPI *VirtualProtectFunc) (LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect);
		typedef BOOL (WINAPI *WriteProcessMemoryFunc) (HANDLE hProcess, LPVOID lpBaseAddress, LPCVOID lpBuffer, SIZE_T nSize, SIZE_T* lpNumberOfBytesWritten);

		GetThreadContextFunc m_GetThreadContext;
		SetThreadContextFunc m_SetThreadContext;
		VirtualProtectFunc m_VirtualProtect;
		WriteProcessMemoryFunc m_WriteProcessMemory;

		DWORD m_MemoryBlockSize;
		DWORD m_CodeMemorySize;
		BYTE* m_pCodeMemory;

		typedef std::vector<CModule*> TModuleVector;
		TModuleVector m_Modules;

		typedef std::vector<CPlugin*> TPluginVector;
		TPluginVector m_Plugins;

		void InitFunctions ();

		CModule* GetModule (LPCWSTR lpModuleName);
};

class CMasterHost : public CHost
{
	public:
		CMasterHost ();
		virtual ~CMasterHost ();

		void UpdateHostState ();

		void AddHost (CHost* pHost);
		void RemoveHost (CHost* pHost);

	protected:
		typedef std::vector<CHost*> THostVector;
		THostVector m_Hosts;
};